{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31c9be60-bb96-4af9-a5ee-1eb47b201d45",
   "metadata": {},
   "source": [
    "# Transformers 量化技术 BitsAndBytes\n",
    "\n",
    "![](docs/images/qlora.png)\n",
    "\n",
    "`bitsandbytes`是将模型量化为8位和4位的最简单选择。 \n",
    "\n",
    "- 8位量化将fp16中的异常值与int8中的非异常值相乘，将非异常值转换回fp16，然后将它们相加以返回fp16中的权重。这减少了异常值对模型性能产生的降级效果。\n",
    "- 4位量化进一步压缩了模型，并且通常与QLoRA一起用于微调量化LLM（低精度语言模型）。\n",
    "\n",
    "（`异常值`是指大于某个阈值的隐藏状态值，这些值是以fp16进行计算的。虽然这些值通常服从正态分布（[-3.5, 3.5]），但对于大型模型来说，该分布可能会有很大差异（[-60, 6]或[6, 60]）。8位量化适用于约为5左右的数值，但超过此范围后将导致显著性能损失。一个好的默认阈值是6，但对于不稳定的模型（小型模型或微调）可能需要更低的阈值。）\n",
    "\n",
    "## 在 Transformers 中使用参数量化\n",
    "\n",
    "使用 Transformers 库的 `model.from_pretrained()`方法中的`load_in_8bit`或`load_in_4bit`参数，便可以对模型进行量化。只要模型支持使用Accelerate加载并包含torch.nn.Linear层，这几乎适用于任何模态的任何模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c2385671-3e67-4fcb-9243-d4b1affea031",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "model_id = \"facebook/opt-2.7b\"\n",
    "\n",
    "model_4bit = AutoModelForCausalLM.from_pretrained(model_id,\n",
    "                                                  device_map=\"auto\",\n",
    "                                                  load_in_4bit=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f4731ac5-fe26-471e-ad17-eb2ba42cb596",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OPTForCausalLM(\n",
       "  (model): OPTModel(\n",
       "    (decoder): OPTDecoder(\n",
       "      (embed_tokens): Embedding(50272, 2560, padding_idx=1)\n",
       "      (embed_positions): OPTLearnedPositionalEmbedding(2050, 2560)\n",
       "      (final_layer_norm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "      (layers): ModuleList(\n",
       "        (0-31): 32 x OPTDecoderLayer(\n",
       "          (self_attn): OPTAttention(\n",
       "            (k_proj): Linear4bit(in_features=2560, out_features=2560, bias=True)\n",
       "            (v_proj): Linear4bit(in_features=2560, out_features=2560, bias=True)\n",
       "            (q_proj): Linear4bit(in_features=2560, out_features=2560, bias=True)\n",
       "            (out_proj): Linear4bit(in_features=2560, out_features=2560, bias=True)\n",
       "          )\n",
       "          (activation_fn): ReLU()\n",
       "          (self_attn_layer_norm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "          (fc1): Linear4bit(in_features=2560, out_features=10240, bias=True)\n",
       "          (fc2): Linear4bit(in_features=10240, out_features=2560, bias=True)\n",
       "          (final_layer_norm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (lm_head): Linear(in_features=2560, out_features=50272, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_4bit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a198b216-b113-4851-a02f-f57be038e1ac",
   "metadata": {},
   "source": [
    "### 实测GPU显存占用：Int4 量化精度\n",
    "\n",
    "```shell\n",
    "Sun Dec 24 18:04:14 2023\n",
    "+---------------------------------------------------------------------------------------+\n",
    "| NVIDIA-SMI 535.129.03             Driver Version: 535.129.03   CUDA Version: 12.2     |\n",
    "|-----------------------------------------+----------------------+----------------------+\n",
    "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
    "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
    "|                                         |                      |               MIG M. |\n",
    "|=========================================+======================+======================|\n",
    "|   0  Tesla T4                       Off | 00000000:00:0D.0 Off |                    0 |\n",
    "| N/A   42C    P0              26W /  70W |   1779MiB / 15360MiB |      0%      Default |\n",
    "|                                         |                      |                  N/A |\n",
    "+-----------------------------------------+----------------------+----------------------+\n",
    "\n",
    "+---------------------------------------------------------------------------------------+\n",
    "| Processes:                                                                            |\n",
    "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
    "|        ID   ID                                                             Usage      |\n",
    "|=======================================================================================|\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1d504b78-9ea4-4100-b614-03dc3bbcb65b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1457.52MiB\n"
     ]
    }
   ],
   "source": [
    "# 获取当前模型占用的 GPU显存（差值为预留给 PyTorch 的显存）\n",
    "memory_footprint_bytes = model_4bit.get_memory_footprint()\n",
    "memory_footprint_mib = memory_footprint_bytes / (1024 ** 2)  # 转换为 MiB\n",
    "\n",
    "print(f\"{memory_footprint_mib:.2f}MiB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8af2edef-9142-443b-b55c-b57872a1fc1d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.11/site-packages/bitsandbytes/nn/modules.py:226: UserWarning: Input type into Linear4bit is torch.float16, but bnb_4bit_compute_type=torch.float32 (default). This will lead to slow inference or training speed.\n",
      "  warnings.warn(f'Input type into Linear4bit is torch.float16, but bnb_4bit_compute_type=torch.float32 (default). This will lead to slow inference or training speed.')\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Merry Christmas! I'm glad to see you're still around.\n",
      "I'm still around, just not posting as much. I'm still here, just not posting as much. I'm still here, just not posting as much. I'm still here, just not posting as much. I'm still here, just not posting as much. I'm\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id\n",
    "                                         )\n",
    "text = \"Merry Christmas! I'm glad to\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\").to(0)\n",
    "\n",
    "out = model_4bit.generate(**inputs, max_new_tokens=64)\n",
    "print(tokenizer.decode(out[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21f299ea-77f6-45cc-82c9-87c96addda06",
   "metadata": {},
   "source": [
    "### 使用 NF4 精度加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "00249404-c60b-47a5-bcb9-a8a4b4b6266f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BitsAndBytesConfig\n",
    "\n",
    "nf4_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    ")\n",
    "\n",
    "model_nf4 = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=nf4_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ab499752-4c53-4ab4-a6a1-1fdf88cbbd0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1457.52MiB\n"
     ]
    }
   ],
   "source": [
    "# 获取当前模型占用的 GPU显存（差值为预留给 PyTorch 的显存）\n",
    "memory_footprint_bytes = model_nf4.get_memory_footprint()\n",
    "memory_footprint_mib = memory_footprint_bytes / (1024 ** 2)  # 转换为 MiB\n",
    "\n",
    "print(f\"{memory_footprint_mib:.2f}MiB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7d335c9-9f13-4834-8008-af20a9f5ca56",
   "metadata": {},
   "source": [
    "### 使用双量化加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f6bfa211-9ad8-4c7b-93a8-37cccaad975a",
   "metadata": {},
   "outputs": [],
   "source": [
    "double_quant_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    ")\n",
    "\n",
    "model_double_quant = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=double_quant_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dbb3913a-a4aa-4d65-8901-8bcf546f1e08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1457.52MiB\n"
     ]
    }
   ],
   "source": [
    "# 获取当前模型占用的 GPU显存（差值为预留给 PyTorch 的显存）\n",
    "memory_footprint_bytes = model_double_quant.get_memory_footprint()\n",
    "memory_footprint_mib = memory_footprint_bytes / (1024 ** 2)  # 转换为 MiB\n",
    "\n",
    "print(f\"{memory_footprint_mib:.2f}MiB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8153e9d-a080-47df-af83-3f1582b2b367",
   "metadata": {},
   "source": [
    "### 使用 QLoRA 所有量化技术加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a4bd4f3a-a7f9-4545-b6a9-732fd6f91b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "qlora_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "model_qlora = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=qlora_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "da8edf77-03cc-4303-a3c0-1b088e5ec958",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1457.52MiB\n"
     ]
    }
   ],
   "source": [
    "# 获取当前模型占用的 GPU显存（差值为预留给 PyTorch 的显存）\n",
    "memory_footprint_bytes = model_qlora.get_memory_footprint()\n",
    "memory_footprint_mib = memory_footprint_bytes / (1024 ** 2)  # 转换为 MiB\n",
    "\n",
    "print(f\"{memory_footprint_mib:.2f}MiB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bf04637-85f6-4ef1-a1bd-81448cd9325c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
